#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/7/25
# @USER    : Shengji He
# @File    : utils.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
import contextlib
import sys

import numpy as np
import torch

import csv

__all__ = [
    'dummy_context', 'empty_cache',
    'nostdout',
    '_normalize',
    'read_csv',
]


def read_csv(file, trace_func=print, encoding='utf-8'):
    out = []
    with open(file, 'r', encoding=encoding) as fr:
        reader = csv.reader(fr)
        for i in reader:
            out.append(i)
    trace_func('Successfully read csv: {}'.format(file))
    return out


class dummy_context(object):
    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


def empty_cache(device: torch.device):
    if device.type == 'cuda':
        torch.cuda.empty_cache()
    elif device.type == 'mps':
        from torch import mps
        mps.empty_cache()
    else:
        pass


"""
Helpers to suppress stdout prints from nnunet
https://stackoverflow.com/questions/2828953/silence-the-stdout-of-a-function-in-python-without-trashing-sys-stdout-and-resto
"""


class DummyFile:
    def write(self, x): pass

    def flush(self): pass


@contextlib.contextmanager
def nostdout(verbose=False):
    if not verbose:
        save_stdout = sys.stdout
        sys.stdout = DummyFile()
        yield
        sys.stdout = save_stdout
    else:
        yield


def CTNormalization(image: np.ndarray, intensityproperties: dict) -> np.ndarray:
    target_dtype = np.float32
    image = image.astype(target_dtype)
    mean_intensity = intensityproperties['mean']
    std_intensity = intensityproperties['std']
    lower_bound = intensityproperties['percentile_00_5']
    upper_bound = intensityproperties['percentile_99_5']
    image = np.clip(image, lower_bound, upper_bound)
    image = (image - mean_intensity) / max(std_intensity, 1e-8)
    return image


def _normalize(data: np.ndarray, intensityproperties: dict) -> np.ndarray:
    for c in range(data.shape[0]):
        data[c] = CTNormalization(data[c], intensityproperties)
    return data
